#include "math.h"
#include "stdlib.h"
#include "inc/core/particles.h"
#include "inc/core/lib.h"
#include "stdio.h"

int compare(void const* one, void const* two)
{
    if (*(floatdig*)one > *(floatdig*)two)
        return 1;
    if (*(floatdig*)one < *(floatdig*)two)
        return -1;
    return 0;
}

__host__ void copyspeedpool(floatdig* d_p, particle* particles, int axe, int shift)
{
	particle* p = particles;
    floatdig* temp = (floatdig*)malloc(sizeof(floatdig)*PARTICLEPOOL);
	int i;
	CUDA_CALL(cudaMemcpy(temp, d_p, sizeof(floatdig)*PARTICLEPOOL, cudaMemcpyDeviceToHost));
	for(i = 0; i < PARTICLEPOOL; i++)
	{
		p[i + shift].speed.v[axe] = temp[i];
	}
	free(temp);
}

__host__ void copypospool(floatdig* d_p, particle* particles, int axe, int shift)
{
	particle* p = particles;
    floatdig* temp = (floatdig*)malloc(sizeof(floatdig)*PARTICLEPOOL);
	int i;
	CUDA_CALL(cudaMemcpy(temp, d_p, sizeof(floatdig)*PARTICLEPOOL, cudaMemcpyDeviceToHost));
	for(i = 0; i < PARTICLEPOOL; i++)
	{
		p[i + shift].pos.v[axe] = temp[i];
	}
	free(temp);
}

void particle_initialize_particle_set(particles_set* set, int N, floatdig charge, floatdig mass)
{
    if (set != NULL)
    {
        if (N > 0)
        {
            particle *particles;
            set->particles = (particle*) malloc(sizeof (particle) * N);
            particles = set->particles;
            if (particles != NULL)
            {
                int i = 0;
                particle* one_particle = NULL;
                for (i = 0; i < N; i++)
                {
                    one_particle = &(particles[i]);
                    lib_set_vector(&(one_particle->pos), 0);
                    lib_set_vector(&(one_particle->speed), 0);
                    one_particle->active = PARTICLE_INACTIVE;
                }
            }
        }
        else
        {
            set->particles = NULL;
        }
        set->num_of_particles = N;
        set->charge = charge;
        set->mass = mass;
        CURAND_CALL(curandCreateGenerator(&set->cuda_r, CURAND_RNG_PSEUDO_DEFAULT));
        CURAND_CALL(curandSetPseudoRandomGeneratorSeed(set->cuda_r, 1234ULL));
		gsl_rng_env_setup();
        set->gsl_r = gsl_rng_alloc(gsl_rng_taus);
		gsl_rng_set (set->gsl_r, 567ULL);
    }
}

__global__ void cpspeed(particle* d_pool, floatdig* d_speed, int axe)
{
    int nparticle = GETPROCPARTICLE(blockDim, blockIdx, threadIdx);
    d_pool[nparticle].speed.v[axe] = d_speed[nparticle];
	d_pool[nparticle].startspeed.v[axe] = d_speed[nparticle];
}

__global__ void cppos(particle* d_pool, floatdig* d_pos, int axe)
{
    int nparticle = GETPROCPARTICLE(blockDim, blockIdx, threadIdx);
    d_pool[nparticle].pos.v[axe] = d_pos[nparticle];
}

__global__ void transform_uniform_distribution(floatdig* d_p, floatdig max, floatdig min)
{
    int n = GETPROCPARTICLE(blockDim, blockIdx, threadIdx);
    d_p[n] = (max - min)*d_p[n] + min;
}


void particle_clear_particle_set(particles_set* set)
{
    if (set != NULL)
    {
        CURAND_CALL(curandDestroyGenerator(set->cuda_r));
		gsl_rng_free(set->gsl_r);
        free(set->particles);
        set->particles = NULL;
        set->num_of_particles = 0;
    }
}

void particle_distribute_particle_set_speed_maxwell_cuda(particles_set* set, floatdig sigma, int axe)
{
	floatdig* d_speed;
    particle* h_p = set->particles;
    int nbunches = set->num_of_particles/PARTICLEPOOL;
	int i;
	#ifdef DEBUG
	printf("particle_distribute_particle_set_speed_maxwell: start: axe %d\n",axe);
	#endif
    CUDA_CALL(cudaMalloc(&d_speed, sizeof(floatdig)*PARTICLEPOOL));
	for (i = 0; i < nbunches; i++)
	{
		CURAND_CALL(curandGenerateNormal(set->cuda_r, d_speed, PARTICLEPOOL, 0, sigma));
		copyspeedpool(d_speed, set->particles, axe, i*PARTICLEPOOL);
	}

    CUDA_CALL(cudaFree(d_speed));
	#ifdef DEBUG
	printf("particle_distribute_particle_set_speed_maxwell: done\n");
	#endif
}

void particle_make_speed_disturbance(particles_set* set, floatdig as, floatdig k, int axe)
{
    int i;
    int N = set->num_of_particles;
    particle* one_particle;
    speed3d* vec;
    for (i = 0; i < N; i++)
    {
        one_particle = &(set->particles[i]);
        vec = &(one_particle->speed);
        vec->v[axe] += as * cos(k * one_particle->pos.v[axe]);
    }
}

void particle_make_pos_disturbance(particles_set* set, floatdig ap, floatdig k, int axe)
{
    int i;
    int N = set->num_of_particles;
    particle* one_particle;
    pos3d* pos;
    for (i = 0; i < N; i++)
    {
        one_particle = &(set->particles[i]);
        pos = &(one_particle->pos);
        pos->v[axe] += ap * sin(k * pos->v[axe]);
    }
}

void particle_distribute_particle_set_beam(particles_set* set, floatdig beam_speed, int axe)
{
    int i;
    int N = set->num_of_particles;
    particle* one_particle;
    speed3d* vec;
    for (i = 0; i < N; i++)
    {
        if (set->particles[i].active == PARTICLE_ACTIVE)
        {
            one_particle = &(set->particles[i]);
            vec = &(one_particle->speed);
            vec->v[axe] = beam_speed;
            one_particle->meanspeed2.v[axe] = beam_speed;
        }

    }
}

void particle_distrubute_particles_set_pos_uniform_cuda(particles_set* set, int axe, floatdig min_pos, floatdig max_pos)
{
    floatdig* d_pos;
    particle* h_p = set->particles;
    int nbunches = set->num_of_particles/PARTICLEPOOL;
	int i;
	#ifdef DEBUG
	printf("particle_distrubute_particles_set_pos_uniform: start: axe %d\n",axe);
	#endif
    CUDA_CALL(cudaMalloc(&d_pos, sizeof(floatdig)*PARTICLEPOOL));
	for (i = 0; i < nbunches; i++)
	{
		CURAND_CALL(curandGenerateUniform(set->cuda_r, d_pos, PARTICLEPOOL));
		transform_uniform_distribution<<<PPOOLSPECIFICATOR>>>(d_pos, max_pos, min_pos);
		copypospool(d_pos, set->particles, axe, i*PARTICLEPOOL);
	}

    CUDA_CALL(cudaFree(d_pos));
	#ifdef DEBUG
	printf("particle_distrubute_particles_set_pos_uniform: done\n");
	#endif
}

particle* particle_get_particle(particles_set* set, int num)
{
    if ((set != NULL) && (set->particles != NULL) && (num < set->num_of_particles))
    {
        return &(set->particles[num]);
    }
    else
    {
        return NULL;
    }
}

void particle_put_particles_pos_into_file(char* fname, particles_set* set, int axe)
{
    FILE* to = fopen(fname, "w");
    if ((to != NULL) && (set != NULL) & (set->particles != NULL))
    {
        int i = 0;
        int N = set->num_of_particles;
        pos3d* pos = NULL;
        for (i = 0; i < N; i++)
        {
            if (set->particles[i].active == PARTICLE_ACTIVE)
            {
                pos = &(set->particles[i].pos);
                fprintf(to, "%d\t%f\n", i, (pos->v[axe]));
            }
        }
        fclose(to);
    }
}

void particle_put_particles_speed_into_file(char* fname, particles_set* set, int axe)
{
    FILE* to = fopen(fname, "w");
    if ((to != NULL) && (set != NULL) & (set->particles != NULL))
    {
        int i = 0;
        int N = set->num_of_particles;
        speed3d* speed = NULL;
        for (i = 0; i < N; i++)
        {
            speed = &(set->particles[i].speed);
            fprintf(to, "%d\t%e\n", i, (speed->v[axe]));
        }
        fclose(to);
    }
}

void particle_add_particle(particles_set* set, particle* one)
{
    if ((set != NULL) && (one != NULL))
    {
        if ((set->particles == NULL) || (set->num_of_particles == 0))
        {
            set->particles = (particle*) malloc(sizeof (particle));
            memcpy(&(set->particles[0]), one, sizeof (particle));
            set->num_of_particles = 1;
        }
        else
        {
            int i = 0;
            particle* new_set = (particle*) malloc(sizeof (particle)*(set->num_of_particles + 1));
            for (i = 0; i < set->num_of_particles; i++)
            {
                memcpy(&(new_set[i]), &(set->particles[i]), sizeof (particle));
            }
            memcpy(&(new_set[set->num_of_particles]), one, sizeof (particle));
            free(set->particles);
            set->particles = new_set;
            set->num_of_particles = set->num_of_particles + 1;
        }
    }
}

__global__ void find_particles_in_pos_range(floatdig* d_zpool, int* d_flags, floatdig pos, floatdig dz)
{
    int nparticle = GETPROCPARTICLE(blockDim, blockIdx, threadIdx);
    d_flags[nparticle] = 0;

    if ( fabs(d_zpool[nparticle] - pos) < dz/2 )
    {
        d_flags[nparticle] = 1;
    }
}

__host__ void copyflagsDeviceToHost(int* d_flags, particle* particles, int shift)
{
	int* temp = (int*)malloc(sizeof(int)*PARTICLEPOOL);
	int i;
	memset(temp, 0, sizeof(int)*PARTICLEPOOL);
	CUDA_CALL(cudaMemcpy(temp, d_flags, sizeof(int)*PARTICLEPOOL, cudaMemcpyDeviceToHost));
	#pragma omp parallel for
	for(i = 0; i < PARTICLEPOOL; i++)
	{
		particles[i + shift].flag = temp[i];
	}
	free(temp);
}

__host__ void copyzpoolDeviceToHost(floatdig* d_zpool, particle* particles, int shift)
{
	floatdig* temp = (floatdig*)malloc(sizeof(floatdig)*PARTICLEPOOL);
	int i;
	memset(temp, 0, sizeof(floatdig)*PARTICLEPOOL);
	CUDA_CALL(cudaMemcpy(temp, d_zpool, sizeof(floatdig)*PARTICLEPOOL, cudaMemcpyDeviceToHost));
	#pragma omp parallel for
	for(i = 0; i < PARTICLEPOOL; i++)
	{
		particles[i + shift].pos.v[Z] = temp[i];
	}
	free(temp);
}

__host__ void copyzpoolHostToDevice(floatdig* d_zpool, particle* particles, int shift)
{
	floatdig* temp = (floatdig*)malloc(sizeof(floatdig)*PARTICLEPOOL);
	int i;
	memset(temp, 0, sizeof(floatdig)*PARTICLEPOOL);
	#pragma omp parallel for
	for(i = 0; i < PARTICLEPOOL; i++)
	{
		temp[i] = particles[i + shift].pos.v[Z];
	}
	CUDA_CALL(cudaMemcpy(d_zpool, temp, sizeof(floatdig)*PARTICLEPOOL, cudaMemcpyHostToDevice));
	free(temp);
}

__host__ void copyflagsHostToDevice(floatdig* d_flags, particle* particles, int shift)
{
	int* temp = (int*)malloc(sizeof(int)*PARTICLEPOOL);
	int i;
	memset(temp, 0, sizeof(int)*PARTICLEPOOL);
	#pragma omp parallel for
	for(i = 0; i < PARTICLEPOOL; i++)
	{
		temp[i] = particles[i + shift].flag;
	}
	CUDA_CALL(cudaMemcpy(d_flags, temp, sizeof(int)*PARTICLEPOOL, cudaMemcpyHostToDevice));
	free(temp);
}

 void particle_dump_particles_pos_into_file_cuda(particles_set* set, char* path, int fid, floatdig pos, floatdig dz)
{
    int i;
    char fullname[MAX_BUF_SIZE];
    FILE* to;
    floatdig* d_zpool;
	int* d_flags;
    particle* p = set->particles;
    int nbunches = set->num_of_particles/PARTICLEPOOL;
    sprintf(fullname, "%s/%s_%d.dat", path, POSFILESNAME, fid);
	int printall = 0;
#ifdef DEBUG
    printf("particle_dump_particles_pos_into_file: Start\nprint to %s\n", fullname);
#endif
    if (dz < 0)
    {
        printall = 1;
        goto PRINT; //oh my god, i used a label!
    }
    CUDA_CALL(cudaMalloc(&d_zpool, sizeof(floatdig)*PARTICLEPOOL));
	CUDA_CALL(cudaMalloc(&d_flags, sizeof(int)*PARTICLEPOOL));

	for(i = 0; i < nbunches; i++)
	{
		copyzpoolHostToDevice(d_zpool, p, i*PARTICLEPOOL);
		find_particles_in_pos_range<<<PPOOLSPECIFICATOR>>>(d_zpool, d_flags, pos, dz);
		copyflagsDeviceToHost(d_flags, p, i*PARTICLEPOOL);
	}
    CUDA_CALL(cudaFree(d_zpool));
	CUDA_CALL(cudaFree(d_flags));
PRINT:
    to = fopen(fullname, "w");
    for (i = 0; i < set->num_of_particles; i++)
    {
        if (printall || set->particles[i].flag)
        {

            fprintf(to, "%e\n", set->particles[i].pos.v[VZ]);
        }
    }
    fclose(to);
#ifdef DEBUG
    printf("particle_dump_particles_pos_into_file: Done\n");
#endif
}


void particle_dump_particles_speed_into_file_cuda(particles_set* set, char* path, int fid, floatdig pos, floatdig dz)
{
    int i;
    char fullname[MAX_BUF_SIZE];
    FILE* to;
    floatdig* d_zpool;
	int* d_flags;
    particle* p = set->particles;
    int nbunches = set->num_of_particles/PARTICLEPOOL;
    DIM3THREADSPERBLOCK;
    sprintf(fullname, "%s/%s_%d.dat", path, SPEEDFILESNAME, fid);
	int printall = 0;
#ifdef DEBUG
    printf("particle_dump_particles_speed_into_file: Start\nprint to %s\n", fullname);
#endif
    if (dz < 0)
    {
        printall = 1;
        goto PRINT; //oh my god, i used a label!
    }
    CUDA_CALL(cudaMalloc(&d_zpool, sizeof(floatdig)*PARTICLEPOOL));
	CUDA_CALL(cudaMalloc(&d_flags, sizeof(int)*PARTICLEPOOL));
	for(i = 0; i < nbunches; i++)
	{
		copyzpoolHostToDevice(d_zpool, p, i*PARTICLEPOOL);
		find_particles_in_pos_range<<<PPOOLSPECIFICATOR>>>(d_zpool, d_flags, pos, dz);
		copyflagsDeviceToHost(d_flags, p, i*PARTICLEPOOL);
	}

    CUDA_CALL(cudaFree(d_zpool));
	CUDA_CALL(cudaFree(d_flags));
PRINT:
    to = fopen(fullname, "w");
    for (i = 0; i < set->num_of_particles; i++)
    {
        if (printall || set->particles[i].flag)
        {

            fprintf(to, "%e\t%e\t%e\t\n", set->particles[i].speed.v[VX], set->particles[i].speed.v[VY], set->particles[i].speed.v[VZ]);
        }
    }
    fclose(to);
#ifdef DEBUG
    printf("particle_dump_particles_speed_into_file: Done\n");
#endif
}

__global__ void gettre(floatdig* tre, particle* d_pool, floatdig m, volatile int* processed)
{
    __shared__ floatdig last;
    int n = blockDim.x*blockIdx.x + threadIdx.x;
	if(d_pool[n].flag)
	{
		tre[*processed] = sqrt(d_pool[n].meanspeed2.v[VX] + d_pool[n].meanspeed2.v[VY]);
		*processed += 1;
	}
}

__global__ void calculate_trenergy_of_flagged_particles()
{

}

void particle_calculate_trenergy_distribution_function_cuda_omp(particles_set* set, char* path, int fid, floatdig pos, floatdig dz, int ensteps)
{
	int i;
    char fullname[MAX_BUF_SIZE];
    FILE* to;
    floatdig* d_zpool;
    particle* p = set->particles;
    int nbunches = set->num_of_particles/PARTICLEPOOL;
    DIM3THREADSPERBLOCK;
	int printall = 0;
	int* d_ncount;
	int ncount;
	int* d_flags;
    int k = 0;
	floatdig* trenergy;
	floatdig tr;
	floatdig* f;
	floatdig* ftrenergy;
    floatdig mintr;
    floatdig maxtr;
    floatdig dtr;
    floatdig Enorm = 0;
    floatdig fnorm = 0;

	sprintf(fullname, "%s/%s_%d.dat", path, TRENFUNCFILESNAME, fid);
#ifdef DEBUG
    printf("particle_calculate_trenergy_distribution_function_cuda_omp: Start\nprint to %s\n", fullname);
#endif
    if (dz < 0)
    {
        printf("particle_calculate_trenergy_distribution_function_cuda_omp: dz should be a positive value\n");

    }
    CUDA_CALL(cudaMalloc(&d_zpool, sizeof(particle)*PARTICLEPOOL));
	CUDA_CALL(cudaMalloc(&d_ncount, sizeof(int)));
	CUDA_CALL(cudaMemset(d_ncount, 0, sizeof(int)));
	CUDA_CALL(cudaMalloc(&d_flags, sizeof(int)*PARTICLEPOOL));
	for(i = 0; i < nbunches; i++)
	{
		copyzpoolHostToDevice(d_zpool, p, i*PARTICLEPOOL);
		find_particles_in_pos_range<<<PPOOLSPECIFICATOR>>>(d_zpool, d_flags, pos, dz);
		copyflagsDeviceToHost(d_flags, p, i*PARTICLEPOOL);
	}

    CUDA_CALL(cudaFree(d_zpool));
	CUDA_CALL(cudaFree(d_flags));
    CUDA_CALL(cudaMemcpy(&ncount, d_ncount, sizeof(int), cudaMemcpyDeviceToHost));
    CUDA_CALL(cudaFree(d_ncount));

	
    //#pragma omp parallel for
    for(i = 0; i < set->num_of_particles; i++)
    {
        if(set->particles[i].flag)
        {
            //trenergy[k] = ELECTRON_MASS*(set->particles[i].meanspeed2.v[VX] + set->particles[i].meanspeed2.v[VY])/2;
			ncount++;
        }
    }
	trenergy = (floatdig*)malloc(ncount*sizeof(floatdig));
    f = (floatdig*)malloc(ensteps*sizeof(floatdig));
	ftrenergy = (floatdig*)malloc(ensteps*sizeof(floatdig));
	memset(f, 0,ensteps*sizeof(floatdig));
	memset(ftrenergy, 0,ensteps*sizeof(floatdig));
	
	//#pragma omp parallel for shared(k)
    for(i = 0; i < set->num_of_particles; i++)
    {
        if(set->particles[i].flag)
        {
            trenergy[k] = ELECTRON_MASS*(set->particles[i].meanspeed2.v[VX] + set->particles[i].meanspeed2.v[VY])/2;
			k++;
        }
    }
	
    //not a parallel sorting :(
    qsort(trenergy, ncount, sizeof(floatdig), compare);
    mintr = trenergy[0];
    maxtr = trenergy[ncount - 1];
    dtr = (maxtr - mintr)/ensteps;

    //#pragma omp parallel for
    for (k = 0; k <ensteps; k++)
    {
        tr = mintr + dtr*k;
		ftrenergy[k] = tr + 0.5*dtr;
        for (i = 0; i < ncount; i++)
        {
            if (fabs(trenergy[i] - tr) < dtr/2)
            {
                f[k]++;
            }
        }
    }
    //#pragma omp parallel for
    for (k = 0; k < ensteps; k++)
    {
        Enorm += f[k]*ftrenergy[k]*dtr;
        fnorm += f[k]*dtr;
    }

    to = fopen(fullname, "w");
    //for (k = 0; k < ensteps; k++)
    //{
    //    fprintf(to, "%e\t%e\n", ftrenergy[k], f[k]);
    //}
    fprintf(to, "%e\n", Enorm/fnorm);
    fclose(to);
	free(trenergy);
	free(f);
	free(ftrenergy);
#ifdef DEBUG
    printf("particle_calculate_trenergy_distribution_function_cuda_omp: Done\n");
#endif

}
